import pandas as pd

from gaussian_process.gaussian_process_regression import GaussianProcessRegressor
from gaussian_process.kernel import RBFKernel
from sklearn.linear_model import LinearRegression
from data_loader import Dataset, Dataset_Pred
import numpy as np
import matplotlib.pyplot as plt
from basis_function import Polynomial


class GPRegressor:
    def __init__(self, model):
        self.model = model

    def fit(self, X, y):
        self.train_x = X
        self.train_y = y
        self.fit(X, y)

    def predict(self, return_cov=True, pred_step=10):
        self.test_x = np.linspace(0, len(self.train_x) + pred_step, 100)
        self.output, self.cov = self.model.predict(self.test_x, return_cov=return_cov)
        return self.output, self.cov

    def sample(self, num=1):
        if hasattr(self.model, 'sample_func'):
            return self.model.sample_func(self.test_x, num)
        else:
            raise NotImplementedError

    def plot(self):
        sample = self.sample()
        uncertainty = 1.96 * np.sqrt(np.diag(self.cov))
        plt.figure(figsize=(20, 5))
        plt.plot(self.test_x, self.output, label='predict', c='green')
        plt.plot(self.test_x, sample.flatten(), label='sample', c='y', linestyle='--')
        plt.scatter(self.train_x, self.train_y, c='red', marker='x', label='observation', linewidths=5)
        plt.fill_between(self.test_x, self.output.flatten() + uncertainty, self.output.flatten() - uncertainty, alpha=0.1)
        plt.legend()
        plt.show()


class PolynomialRegressor:
    def __init__(self, degree):
        self.basis_func = Polynomial(degree)
        self.model = LinearRegression()

    def fit(self, X, y):
        self.train_x = X
        self.train_y = y
        self.model.fit(self.basis_func(X), y)

    def predict(self, X):
        self.test_x = X
        predict = self.model.predict(self.basis_func(X))
        self.output = predict
        return predict

    def plot(self):
        plt.figure(figsize=(20, 5))
        plt.plot(self.test_x, self.output, label='predict', c='green')
        plt.plot(self.train_x, self.train_y, c='red', label='observation')
        plt.ylim(-5, 20)
        plt.legend()
        plt.show()


if __name__ == '__main__':
    dataset = Dataset_Pred()
    model = GaussianProcessRegressor(kernel=RBFKernel, sigma_overall=0.01, l_scale=2, alpha=1e-4)
    model.fit(dataset.train_data, dataset.target)
    pred, cov = model.predict(dataset.data[:, :-1])
    uncertainty = 1.96 * np.sqrt(np.diag(cov))
    samples = model.sample_func(dataset.pred_data[:, -1])
    result = np.concatenate([dataset.data[:, :-1, ], pred], axis=-1)
    result = dataset.inverse_transform(result)
    frame = dataset.raw_data
    frame['中国油价'] = result[:, -1:]
    pd.DataFrame.to_csv(frame, index=False, path_or_buf='./results/中国油价/pred_results.csv')

    plt.figure(figsize=(15, 5))
    plt.scatter(np.arange(len(dataset.target)), dataset.target, c='red', marker='x', label='observation', linewidths=5)
    plt.plot(pred, label='predict')
    plt.fill_between(np.arange(len(dataset.data)), pred.flatten()-uncertainty, pred.flatten()+uncertainty, alpha=0.3, color='yellow')
    plt.legend()
    plt.savefig('./results/中国油价/pred_results(normalized).pdf')
    plt.show()
